package stone926.mods.more_enchantments.spawn;

import net.minecraft.entity.EntityType;
import net.minecraft.entity.mob.MobEntity;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

public class SimplePool {

  private final List<Entry<? extends MobEntity>> entries = new ArrayList<>();
  private int totalWeight = 0;

  public SimplePool() { }

  public SimplePool(SimplePool pool) {
    putAll(pool.getEntries());
  }

  public static SimplePool of(Entry<?>... entries) {
    SimplePool pool = new SimplePool();
    pool.putAll(Arrays.stream(entries).toList());
    return pool;
  }

  public List<Entry<? extends MobEntity>> getEntries() {
    return entries;
  }

  public int getTotalWeight() {
    return totalWeight;
  }

  public SimplePool put(Entry<? extends MobEntity> entry) {
    entries.add(entry);
    totalWeight += entry.weight();
    return this;
  }

  public SimplePool putAll(List<Entry<? extends MobEntity>> list) {
    list.forEach(this::put);
    return this;
  }

  public SimplePool putAll(Entry<?>... list) {
    return putAll(Arrays.stream(list).toList());
  }

  public Entry<? extends MobEntity> getRandomEntry(Random random) {
    // 50 1 10 // r = 60
    if (totalWeight > 0) {
      int r = random.nextInt(totalWeight + 1);
      int w = 0;
      for (Entry<? extends MobEntity> e : entries) {
        int weight = e.weight();
        if (w <= r && r <= w + weight) {
          return e;
        }
        w += weight;
      }
    }
    return null;
  }

  @Override
  public String toString() {
    StringBuilder entries = new StringBuilder();
    for (Entry<? extends MobEntity> entry : this.entries) {
      entries.append(" ").append(entry).append("\n");
    }
    return "SimplePool{\n" +
      "entries=" + entries +
      ",\ntotalWeight=" + totalWeight +
      "}\n";
  }

  public record Entry<T extends MobEntity>(
    EntityType<T> entityType, EntityInitializer<? super T> entityInitializer, int weight
  ) {

    public static <U extends MobEntity> Entry<U> of(EntityType<U> entityType, EntityInitializer<? super U> entityInitializer, int weight) {
      return new Entry<>(entityType, entityInitializer, weight);
    }

    @Override
    public String toString() {
      return "Entry{" +
        "entityType=" + entityType +
        ", weight=" + weight +
        '}';
    }

  }

}
